#include "include/TcpServer.h"
#include <string>
#include <thread>
#include <mutex>
#include <unordered_map>
#include <atomic>
#include <cstring>
#include <vector>
#include "include/ILogger.h"
#include <unordered_map>

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#include <mswsock.h>
#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib, "mswsock.lib")
#else
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/epoll.h>
#include <errno.h>
#endif

namespace Tool
{

#define BUFFER_SIZE 4096

// Add cross-platform socket definitions
#ifndef _WIN32
typedef int SOCKET;
#define INVALID_SOCKET -1
#define SOCKET_ERROR -1
#endif

struct TcpServer::MessageBuffer 
{
    std::vector<uint8_t> buffer;
    size_t expectedLength = 0;
    size_t receivedLength = 0;
    bool lengthReceived = false;

    void Reset() 
    {
        buffer.clear();
        expectedLength = 0;
        receivedLength = 0;
        lengthReceived = false;
    }
};


// Move ITcpServerPrivater definition here before it's used
struct TcpServer::ITcpServerPrivater
{
    ITcpListener* m_listener;
    std::atomic<bool> m_running{false};

    SOCKET m_listenSocket = INVALID_SOCKET;
    #ifdef _WIN32
        HANDLE m_iocp = INVALID_HANDLE_VALUE;
    #else
        int m_epollFd = -1;
    #endif
    std::unordered_map<std::string, SOCKET> m_clients;
    std::mutex m_clientsMutex;
    int m_threadCount;
    std::vector<std::thread> m_threads;
    std::unordered_map<std::string, MessageBuffer> m_clientBuffers;
};

enum class TcpServer::IOType
{
    Read,
    Write
};

struct TcpServer::IOContext
{
#ifdef _WIN32
    OVERLAPPED overlapped{};
    WSABUF wsaBuf{BUFFER_SIZE, buffer};
    CHAR buffer[BUFFER_SIZE]{};
#else
    char buffer[BUFFER_SIZE]{};
#endif
    IOType type{};
    SOCKET socket = INVALID_SOCKET;
#ifdef _WIN32
    DWORD nBytes = 0;
#else
    ssize_t nBytes = 0;
#endif
};

TcpServer::TcpServer(ITcpListener *listener, int threadCount)
    : m_privater(std::make_unique<ITcpServerPrivater>())
{
    m_privater->m_listener = listener;
    m_privater->m_threadCount = threadCount;
}

TcpServer::~TcpServer()
{
    Stop();
}

bool TcpServer::Start(const char *bindAddress, uint16_t port)
{
#ifdef _WIN32
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        LOG_ERROR("WSAStartup failed: %d", WSAGetLastError());
        return false;
    }

    unsigned long ul = 1;
    m_privater->m_listenSocket = WSASocketW(AF_INET, SOCK_STREAM, 0, nullptr, 0, WSA_FLAG_OVERLAPPED);
#else
    // Linux epoll implementation
    m_privater->m_listenSocket = socket(AF_INET, SOCK_STREAM, 0);
    int opt = 1;
    setsockopt(m_privater->m_listenSocket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
#endif
    if (INVALID_SOCKET == m_privater->m_listenSocket)
    {
#ifdef _WIN32
        LOG_ERROR("FAILED TO CREATE SERVER SOCKET: %d", WSAGetLastError());
#else
        LOG_ERROR("FAILED TO CREATE SERVER SOCKET: %d", errno);
#endif
        return false;
    }

    // 设置为非阻塞 IO
#ifdef _WIN32
    if (SOCKET_ERROR == ioctlsocket(m_privater->m_listenSocket, FIONBIO, &ul))
#else
    int flags = fcntl(m_privater->m_listenSocket, F_GETFL, 0);
    if (fcntl(m_privater->m_listenSocket, F_SETFL, flags | O_NONBLOCK) == -1)
#endif
    {
#ifdef _WIN32
        LOG_ERROR("FAILED TO SET NONBLOCKING SOCKET: %d", WSAGetLastError());
#else
        LOG_ERROR("FAILED TO SET NONBLOCKING SOCKET: %d", errno);
#endif
        return false;
    }

    sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(m_privater->m_listenSocket, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)
    {
#ifdef _WIN32
        LOG_ERROR("Bind failed: %d", WSAGetLastError());
#else
        LOG_ERROR("Bind failed: %d", errno);
#endif
        return false;
    }

    if (listen(m_privater->m_listenSocket, SOMAXCONN) == SOCKET_ERROR)
    {
#ifdef _WIN32
        LOG_ERROR("Listen failed: %d", WSAGetLastError());
#else
        LOG_ERROR("Listen failed: %d", errno);
#endif
        return false;
    }

#ifdef _WIN32
    m_privater->m_iocp = CreateIoCompletionPort(INVALID_HANDLE_VALUE, nullptr, 0, m_privater->m_threadCount);
    if (!m_privater->m_iocp)
    {
        LOG_ERROR("CreateIoCompletionPort failed: %d", GetLastError());
        return false;
    }

    for (int i = 0; i < m_privater->m_threadCount; ++i)
    {
        m_privater->m_threads.emplace_back(&TcpServer::EventWorkerThread, this);
    }
#else
    // Linux epoll setup
    m_privater->m_epollFd = epoll_create1(0);
    if (m_privater->m_epollFd == -1)
    {
        LOG_ERROR("epoll_create failed: %d", errno);
        return false;
    }

    struct epoll_event ev;
    ev.events = EPOLLIN;
    ev.data.fd = m_privater->m_listenSocket;
    if (epoll_ctl(m_privater->m_epollFd, EPOLL_CTL_ADD, m_privater->m_listenSocket, &ev) == -1)
    {
        LOG_ERROR("epoll_ctl failed: %d", errno);
        return false;
    }

    // Start epoll worker thread
    for (int i = 0; i < m_privater->m_threadCount; ++i)
    {
        m_privater->m_threads.emplace_back(&TcpServer::EpollWorkerThread, this);
    }
#endif

    m_privater->m_running = true;
    std::thread([this]() { AcceptLoop(); }).detach();

    LOG_INFO("Server started on port %d", port);
    return true;
}

bool TcpServer::Stop()
{
    m_privater->m_running = false;

    if (m_privater->m_listenSocket != INVALID_SOCKET)
    {
#ifdef _WIN32
        closesocket(m_privater->m_listenSocket);
#else
        close(m_privater->m_listenSocket);
#endif
        m_privater->m_listenSocket = INVALID_SOCKET;
    }

#ifdef _WIN32
    if (m_privater->m_iocp != INVALID_HANDLE_VALUE)
    {
        CloseHandle(m_privater->m_iocp);
        m_privater->m_iocp = INVALID_HANDLE_VALUE;
    }
#endif

    for (auto& thread : m_privater->m_threads)
    {
        if (thread.joinable())
        {
            thread.join();
        }
    }

#ifdef _WIN32
    WSACleanup();
#endif

    return true;
}

bool TcpServer::Send(const char *address, uint16_t port, const void *buffer, size_t length)
{
    std::lock_guard<std::mutex> lock(m_privater->m_clientsMutex);

    std::string key = std::string(address) + ":" + std::to_string(port);
    auto it = m_privater->m_clients.find(key);
    if(it == m_privater->m_clients.end())
    {
        LOG_WARN("Send failed: client not found - %s", key.c_str());
        return false;
    }

#ifdef _WIN32
    SOCKET clientSocket = it->second;
    size_t totalLength = sizeof(uint32_t) + length;
    IOContext* ioContext = new IOContext();
    ZeroMemory(&ioContext->overlapped, sizeof(OVERLAPPED));
    ioContext->socket = clientSocket;
    ioContext->type = IOType::Write;
    ioContext->wsaBuf.len = static_cast<ULONG>(totalLength);
    ioContext->wsaBuf.buf = new char[totalLength];

    // 将长度写入前 4 字节（网络字节序）
    uint32_t networkLength = htonl(static_cast<uint32_t>(length));
    memcpy(ioContext->wsaBuf.buf, &networkLength, sizeof(networkLength));
    memcpy(ioContext->wsaBuf.buf + sizeof(networkLength), buffer, length);

    DWORD bytesSent = 0;
    int result = WSASend(clientSocket, &ioContext->wsaBuf, 1, &bytesSent, 0, &ioContext->overlapped, nullptr);
    if (result == SOCKET_ERROR)
    {
        int error = WSAGetLastError();
        if (error != WSA_IO_PENDING)
        {
            LOG_ERROR("WSASend failed: %d", error);
            delete[] ioContext->wsaBuf.buf;
            delete ioContext;
            return false;
        }
    }
#else
    SOCKET clientSocket = it->second;
    size_t totalSent = 0;

    uint32_t networkLength = htonl(static_cast<uint32_t>(length));
    std::vector<char> totalBuffer(sizeof(networkLength) + length);
    memcpy(totalBuffer.data(), &networkLength, sizeof(networkLength));
    memcpy(totalBuffer.data() + sizeof(networkLength), buffer, length);

    while (totalSent < totalBuffer.size())
    {
        ssize_t sent = send(clientSocket, totalBuffer.data() + totalSent, totalBuffer.size() - totalSent, 0);
        if (sent == -1)
        {
            if (errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            LOG_ERROR("send failed: %d", errno);
            return false;
        }
        totalSent += sent;
    }
#endif
    return true;
}

bool TcpServer::Disconnect(const char *address, uint16_t port)
{
    std::string key = std::string(address) + ":" + std::to_string(port);
    auto it = m_privater->m_clients.find(key);
    if(it == m_privater->m_clients.end())
    {
        LOG_WARN("Send failed: client not found - %s", key.c_str());
        return false;
    }
#ifdef _WIN32
    closesocket(it->second);
#endif
    m_privater->m_clients.erase(it);
    return true;
}

#ifdef _WIN32
void TcpServer::EventWorkerThread()
{
    IOContext *ioContext = nullptr;
    DWORD lpNumberOfBytesTransferred = 0;
    void *lpCompletionKey = nullptr;
    DWORD dwFlags = 0;
    DWORD nBytes = BUFFER_SIZE;

    while (m_privater->m_running)
    {
        if (!GetQueuedCompletionStatus(m_privater->m_iocp, &lpNumberOfBytesTransferred, (PULONG_PTR) &lpCompletionKey,(LPOVERLAPPED *) &ioContext, INFINITE))
        {
            if (!m_privater->m_running)
            {
                break;
            }
            LOG_ERROR("GetQueuedCompletionStatus failed: %d", GetLastError());
            continue;
        }

        // 收到 PostQueuedCompletionStatus 发出的退出指令
        if (lpNumberOfBytesTransferred == -1)
        {
            break;
        }

        // 处理客户端正常断开连接（接收到0字节）
        if (lpNumberOfBytesTransferred == 0)
        {
            sockaddr_in clientAddr;
            int addrLen = sizeof(clientAddr);
            if (getpeername(ioContext->socket, (sockaddr*)&clientAddr, &addrLen) == 0)
            {
                char ip[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, &clientAddr.sin_addr, ip, INET_ADDRSTRLEN);
                uint16_t port = ntohs(clientAddr.sin_port);
                m_privater->m_listener->OnDisconnect(ip, port);
                std::lock_guard<std::mutex> lock(m_privater->m_clientsMutex);
                std::string key = std::string(ip) + ":" + std::to_string(port);
                m_privater->m_clients.erase(key);
                LOG_INFO("Client disconnected: %s", key.c_str());
            }
            closesocket(ioContext->socket);
            delete ioContext;
            ioContext = nullptr;
            continue;
        }

        // 读到，或者写入的字节总数
        ioContext->nBytes = lpNumberOfBytesTransferred;
        switch (ioContext->type)
        {
        case IOType::Read:
        {
            // 通知监听器有数据接收
            sockaddr_in clientAddr;
            int addrLen = sizeof(clientAddr);
            if (getpeername(ioContext->socket, (sockaddr*)&clientAddr, &addrLen) == 0)
            {
                char ip[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, &clientAddr.sin_addr, ip, INET_ADDRSTRLEN);
                uint16_t port = ntohs(clientAddr.sin_port);

                // std::lock_guard<std::mutex> lock(m_privater->m_clientsMutex);
                std::string key = std::string(ip) + ":" + std::to_string(port);
                auto it = m_privater->m_clientBuffers.find(key);
                if (it == m_privater->m_clientBuffers.end())
                {
                    m_privater->m_clientBuffers[key] = MessageBuffer{};
                    it = m_privater->m_clientBuffers.find(key);
                }

                // 将新数据追加到缓冲区
                MessageBuffer& buffer = it->second;
                buffer.buffer.insert(buffer.buffer.end(), ioContext->buffer, ioContext->buffer + lpNumberOfBytesTransferred);
                buffer.receivedLength += lpNumberOfBytesTransferred;

                if (!buffer.lengthReceived)
                {
                    if (buffer.receivedLength >= sizeof(uint32_t))
                    {
                        // 解析长度字段（网络字节序转主机字节序）
                        uint32_t networkLength = *reinterpret_cast<uint32_t*>(buffer.buffer.data());
                        buffer.expectedLength = ntohl(networkLength);
                        buffer.lengthReceived = true;
                    }
                    else
                    {
                        // 数据不足，继续接收
                        goto ContinueRead;
                    }
                }

                if (buffer.receivedLength >= buffer.expectedLength + sizeof(uint32_t))
                {
                    buffer.buffer.erase(buffer.buffer.begin(), buffer.buffer.begin() + 4);
                    m_privater->m_listener->OnReceive(ip, port, buffer.buffer.data(), buffer.expectedLength);
                    buffer.Reset();
                }
            }

            // 读取到的字节总数等于最大缓存大小，可能还有内容没有读完，再次提交读取事件
            // 无论接收多少数据，始终继续投递接收请求
ContinueRead:
            int nRt = WSARecv(ioContext->socket, &ioContext->wsaBuf, 1, &nBytes, &dwFlags, &(ioContext->overlapped), nullptr);
            auto e = WSAGetLastError();
            if (SOCKET_ERROR == nRt && e != WSA_IO_PENDING)
            {
                // 获取客户端地址信息
                sockaddr_in clientAddr;
                int addrLen = sizeof(clientAddr);
                if (getpeername(ioContext->socket, (sockaddr*)&clientAddr, &addrLen) == 0)
                {
                    char ip[INET_ADDRSTRLEN];
                    inet_ntop(AF_INET, &clientAddr.sin_addr, ip, INET_ADDRSTRLEN);
                    uint16_t port = ntohs(clientAddr.sin_port);
                    // 通知监听器客户端断开
                    m_privater->m_listener->OnDisconnect(ip, port);
                    // 从客户端列表移除
                    std::lock_guard<std::mutex> lock(m_privater->m_clientsMutex);
                    std::string key = std::string(ip) + ":" + std::to_string(port);
                    m_privater->m_clients.erase(key);
                    LOG_INFO("Client disconnected: %s", key.c_str());
                }
                closesocket(ioContext->socket);
                delete ioContext;
                ioContext = nullptr;
            }
            break;
        }
        case IOType::Write:
        {
            // 释放发送缓冲区和IO上下文
            delete[] ioContext->wsaBuf.buf;
            delete ioContext;
            ioContext = nullptr;
            break;
        }
        }
    }
}
#else
void TcpServer::EpollWorkerThread()
{
    struct epoll_event events[1024];
    while (m_privater->m_running)
    {
        int nfds = epoll_wait(m_privater->m_epollFd, events, 1024, -1);
        if (nfds == -1)
        {
            if (errno == EINTR)
                continue;
            LOG_ERROR("epoll_wait failed: %d", errno);
            break;
        }

        for (int i = 0; i < nfds; ++i)
        {
            if (events[i].data.fd == m_privater->m_listenSocket)
            {
                // New connection, handled in AcceptLoop
                continue;
            }

            if (events[i].events & EPOLLIN)
            {
                HandleReadEvent(events[i].data.fd);
            }
            else if (events[i].events & EPOLLOUT)
            {
                HandleWriteEvent(events[i].data.fd);
            }
            else if (events[i].events & (EPOLLERR | EPOLLHUP))
            {
                HandleErrorEvent(events[i].data.fd);
            }
        }
    }
}

void TcpServer::HandleReadEvent(int clientSocket)
{
    char buffer[BUFFER_SIZE];
    sockaddr_in clientAddr;
    socklen_t addrLen = sizeof(clientAddr);
    getpeername(clientSocket, (struct sockaddr*)&clientAddr, &addrLen);
    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &clientAddr.sin_addr, ip, INET_ADDRSTRLEN);
    uint16_t port = ntohs(clientAddr.sin_port);
    std::string key = std::string(ip) + ":" + std::to_string(port);

    auto it = m_privater->m_clientBuffers.find(key);
    if (it == m_privater->m_clientBuffers.end())
    {
        m_privater->m_clientBuffers[key] = MessageBuffer{};
        it = m_privater->m_clientBuffers.find(key);
    }

    MessageBuffer& msgBuffer = it->second;
    while (true) // 循环读取数据
    {
        ssize_t bytesRead = recv(clientSocket, buffer, BUFFER_SIZE, 0);
        if (bytesRead <= 0)
        {
            if (bytesRead < 0 && (errno == EAGAIN || errno == EWOULDBLOCK))
            {
                return; // 非阻塞模式下无数据可读
            }

            std::lock_guard<std::mutex> lock(m_privater->m_clientsMutex);
            for (auto it = m_privater->m_clients.begin(); it != m_privater->m_clients.end(); )
            {
                if (it->second == clientSocket)
                {
                    m_privater->m_listener->OnDisconnect(it->first.c_str(), 0);
                    it = m_privater->m_clients.erase(it);
                }
                else
                {
                    ++it;
                }
            }
            close(clientSocket);
            return;
        }

        msgBuffer.buffer.insert(msgBuffer.buffer.end(), buffer, buffer + bytesRead);
        msgBuffer.receivedLength += bytesRead;

        if (!msgBuffer.lengthReceived && msgBuffer.receivedLength >= sizeof(uint32_t))
        {
            uint32_t networkLength = *reinterpret_cast<uint32_t*>(msgBuffer.buffer.data());
            msgBuffer.expectedLength = ntohl(networkLength);
            msgBuffer.lengthReceived = true;
        }
        // LOG_INFO("Read size: %d/%d", msgBuffer.receivedLength, msgBuffer.expectedLength);
        if (msgBuffer.receivedLength >= msgBuffer.expectedLength + sizeof(uint32_t))
        {
            msgBuffer.buffer.erase(msgBuffer.buffer.begin(), msgBuffer.buffer.begin() + sizeof(uint32_t));
            m_privater->m_listener->OnReceive(ip, port, msgBuffer.buffer.data(), msgBuffer.expectedLength);
            msgBuffer.Reset();
            break;
        }
    }
}

void TcpServer::HandleWriteEvent(int clientSocket)
{
    // 实现写事件处理逻辑（如需要）
}

void TcpServer::HandleErrorEvent(int clientSocket)
{
    LOG_ERROR("Socket error occurred: %d", errno);
    std::lock_guard<std::mutex> lock(m_privater->m_clientsMutex);
    for (auto it = m_privater->m_clients.begin(); it != m_privater->m_clients.end(); )
    {
        if (it->second == clientSocket)
        {
            m_privater->m_listener->OnDisconnect(it->first.c_str(), 0);
            it = m_privater->m_clients.erase(it);
        }
        else
        {
            ++it;
        }
    }
    close(clientSocket);
}
#endif

// 将AcceptLoop定义移到_WIN32条件块外，并添加Linux实现
void TcpServer::AcceptLoop()
{
#ifdef _WIN32
    while (m_privater->m_running)
    {
        SOCKET clientSocket = accept(m_privater->m_listenSocket, NULL, NULL);
        if (clientSocket == INVALID_SOCKET)
        {
            if (!m_privater->m_running)
            {
                break;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            continue;
        }

        // 将接入的套接字也设置为非阻塞的
        unsigned long ul = 1;
        if (SOCKET_ERROR == ioctlsocket(clientSocket, FIONBIO, &ul))
        {
            shutdown(clientSocket, SD_BOTH);
            closesocket(clientSocket);
            continue;
        }

        if (!CreateIoCompletionPort((HANDLE)clientSocket, m_privater->m_iocp, 0, 0))
        {
            LOG_ERROR("CreateIoCompletionPort for client failed: %d", GetLastError());
            shutdown(clientSocket, SD_BOTH);
            closesocket(clientSocket);
            continue;
        }

        DWORD nBytes = BUFFER_SIZE;
        DWORD dwFlags = 0;
        auto ioContext = new IOContext;
        ioContext->socket = clientSocket;
        ioContext->type = IOType::Read;
        // 此处的 WSARecv 应当理解为 PostRecv，提交一次读取事件
        auto rt = WSARecv(clientSocket, &ioContext->wsaBuf, 1, &nBytes, &dwFlags, &ioContext->overlapped, nullptr);
        auto err = WSAGetLastError();
        if (SOCKET_ERROR == rt && ERROR_IO_PENDING != err)
        {
            // 发生不为 ERROR_IO_PENDING 的错误
            shutdown(clientSocket, SD_BOTH);
            closesocket(clientSocket);
            delete ioContext;
            continue;
        }

        {
            std::lock_guard<std::mutex> lock(m_privater->m_clientsMutex);
            sockaddr_in clientAddr;
            int addrLen = sizeof(clientAddr);
            if (getpeername(clientSocket, (sockaddr*)&clientAddr, &addrLen) == 0)
            {
                char ip[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, &clientAddr.sin_addr, ip, INET_ADDRSTRLEN);
                uint16_t port = ntohs(clientAddr.sin_port);
                std::string key = std::string(ip) + ":" + std::to_string(port);
                m_privater->m_clients[key] = clientSocket;
                m_privater->m_listener->OnConnect(ip, port); // 通知监听器有新连接
                LOG_INFO("Client connected: %s", key.c_str());
            }
        }
    }
#else
    while (m_privater->m_running)
    {
        sockaddr_in clientAddr;
        socklen_t clientAddrLen = sizeof(clientAddr);
        SOCKET clientSocket = accept(m_privater->m_listenSocket, (struct sockaddr*)&clientAddr, &clientAddrLen);
        if (clientSocket == -1)
        {
            if (errno != EAGAIN && errno != EWOULDBLOCK)
            {
                LOG_ERROR("accept failed: %d", errno);
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            continue;
        }

        // 设置客户端套接字为非阻塞
        int flags = fcntl(clientSocket, F_GETFL, 0);
        if (fcntl(clientSocket, F_SETFL, flags | O_NONBLOCK) == -1)
        {
            LOG_ERROR("fcntl failed: %d", errno);
            close(clientSocket);
            continue;
        }

        // 添加到epoll
        struct epoll_event ev;
        ev.events = EPOLLIN | EPOLLET;
        ev.data.fd = clientSocket;
        if (epoll_ctl(m_privater->m_epollFd, EPOLL_CTL_ADD, clientSocket, &ev) == -1)
        {
            LOG_ERROR("epoll_ctl failed: %d", errno);
            close(clientSocket);
            continue;
        }

        // 记录客户端信息
        char ip[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &clientAddr.sin_addr, ip, INET_ADDRSTRLEN);
        uint16_t port = ntohs(clientAddr.sin_port);
        std::string key = std::string(ip) + ":" + std::to_string(port);
        {
            std::lock_guard<std::mutex> lock(m_privater->m_clientsMutex);
            m_privater->m_clients[key] = clientSocket;
        }

        m_privater->m_listener->OnConnect(ip, port);
        LOG_INFO("Client connected: %s", key.c_str());
    }
#endif
}

} // namespace Tool
